# -*- coding:utf-8 -*-
import math
from collections import defaultdict


class TFIDF(object):
    def __init__(self, idf_path):
        self._idf_dict = self.load_idf_dict(idf_path)
        pass

    @staticmethod
    def load_idf_dict(idf_path):
        idf_dict = dict()
        with open(idf_path, 'r', encoding='utf-8') as fr:
            for i, line in enumerate(fr):
                if i == 0:
                    continue
                ss = line.rstrip().split('\t')
                assert len(ss) == 2
                idf_dict[ss[0]] = float(ss[1])
        return idf_dict

    def get_word_idf(self, word):
        return self._idf_dict[word] if word in self._idf_dict else 0.0

    def score(self, query: list, candidate: list):
        d_sent1, d_sent2 = defaultdict(float), defaultdict(float)

        for word in query:
            d_sent1[word] += self.get_word_idf(word)

        for word in candidate:
            d_sent2[word] += self.get_word_idf(word)

        qq_vec, cc_vec, qc_vec = 0.0, 0.0, 0.0
        for word in d_sent1:
            qq_vec += d_sent1[word] ** 2
        for word in d_sent2:
            cc_vec += d_sent2[word] ** 2
        for word in d_sent1:
            if word in d_sent2:
                qc_vec += d_sent1[word] ** 2
        if qq_vec == 0.0 or cc_vec == 0.0:
            return 0.0
        return qc_vec / math.sqrt(qq_vec) / math.sqrt(cc_vec)
